#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <sys/types.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>
#include <tss2/tss2_tpm2_types.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

TSS2_TCTI_CONTEXT *tcti_context = NULL;
TSS2_SYS_CONTEXT *sys_context = NULL;

// Record handles for cleanup
typedef struct HandleRecorder {
    TPM2_HANDLE handles[128];
    uint16_t count;
} HandleRecorder;

static HandleRecorder handle_recorder = {
    .count = 0, .handles = {TPM2_RH_NULL} // Initialize with NULL handle
};

#define CHECK_RETVAL(rc, expected, prefix)                                     \
    {                                                                          \
        if ((rc) != expected) {                                                \
            printf("%sFAILED [%s:%u] %s\n", prefix, __FUNCTION__, __LINE__,    \
                   Tss2_RC_Decode(rc));                                        \
            cleanup(rc);                                                       \
        } else {                                                               \
            printf("%sPASSED [%s:%u]\n", prefix, __FUNCTION__, __LINE__);      \
        }                                                                      \
    }

// Check Passed Case
#define CHECK_PASSED(rc) CHECK_RETVAL(rc, TPM2_RC_SUCCESS, "PassedCase: ")
// Check Failed Case
#define CHECK_FAILED(rc, expected) CHECK_RETVAL(rc, expected, "FailedCase: ")
// Check if the return code indicates success, clean up if not
#define CHECK_RC(rc) CHECK_RETVAL(rc, TPM2_RC_SUCCESS, "")

// Initialize TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context) {
    TSS2_RC rc = Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
    return rc;
}

// Initialize SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context,
                         TSS2_SYS_CONTEXT **sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT *)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size,
                                     tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }

    return rc;
}

// Startup TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc =
        Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_TPM_PROPERTIES,
                               TPM2_PT_MANUFACTURER, 1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

// Cleanup resources
static void cleanup(int rc) {
    for (uint16_t i = 0; i < handle_recorder.count; i++) {
        if (handle_recorder.handles[i] != TPM2_RH_NULL) {
            TSS2_RC rc_flush =
                Tss2_Sys_FlushContext(sys_context, handle_recorder.handles[i]);
            if (rc_flush != TSS2_RC_SUCCESS) {
                printf("FlushContext Failed for handle 0x%08x\n",
                       handle_recorder.handles[i]);
            }
        }
    }
    handle_recorder.count = 0;
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
        free(tcti_context);
    }
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    exit(rc != TSS2_RC_SUCCESS);
    return;
}

#define NV_INDEX_TEST 0x01500015
#define INIT_SIMPLE_TPM2B_SIZE(type) (type).size = sizeof(type) - 2

TSS2_RC nv_define(TSS2_SYS_CONTEXT *sys_context, TPMA_NV attributes,
                  TPMI_RH_PROVISION auth_handle, TPM2B_AUTH *hi_auth,
                  TPMI_RH_NV_INDEX nv_index, TPM2B_AUTH *nv_auth,
                  UINT16 data_size) {
    TSS2_RC rc;
    TSS2L_SYS_AUTH_COMMAND cmd_auths = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = hi_auth->size},
        }
    };
    memcpy(cmd_auths.auths[0].hmac.buffer, hi_auth->buffer, hi_auth->size);
    TSS2L_SYS_AUTH_RESPONSE rsp_auths = {0};
    TPM2B_NV_PUBLIC in_public = {
        .size = 0,
        .nvPublic =
            {
                .nvIndex = nv_index,
                .nameAlg = TPM2_ALG_SHA256,
                .attributes = attributes,
                .authPolicy = {.size = 0},
                .dataSize = data_size,
            },
    };
    rc = Tss2_Sys_NV_DefineSpace(sys_context, auth_handle, &cmd_auths, nv_auth,
                                 &in_public, &rsp_auths);
    return rc;
}

TSS2_RC nv_read(TSS2_SYS_CONTEXT *sys_context, TPMI_RH_NV_INDEX nv_index,
                TPMI_RH_NV_AUTH auth_handle, TPM2B_AUTH *nv_auth,
                TPM2B_MAX_NV_BUFFER *nv_data_read, UINT16 data_size) {
    TSS2_RC rc;
    TSS2L_SYS_AUTH_COMMAND cmd_auths = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = nv_auth->size},
        }
    };
    memcpy(cmd_auths.auths[0].hmac.buffer, nv_auth->buffer, nv_auth->size);
    TSS2L_SYS_AUTH_RESPONSE rsp_auths = {0};

    INIT_SIMPLE_TPM2B_SIZE(*nv_data_read);
    rc = Tss2_Sys_NV_Read(sys_context, auth_handle, nv_index, &cmd_auths,
                          data_size, 0, nv_data_read, &rsp_auths);
    return rc;
}

TSS2_RC nv_write(TSS2_SYS_CONTEXT *sys_context, TPMI_RH_NV_INDEX nv_index,
                 TPMI_RH_NV_AUTH auth_handle, TPM2B_AUTH *nv_auth,
                 TPM2B_MAX_NV_BUFFER *nv_data_write) {
    TSS2_RC rc;
    TSS2L_SYS_AUTH_COMMAND cmd_auths = {
        .count = 1,          
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = nv_auth->size},
        }
    };
    memcpy(cmd_auths.auths[0].hmac.buffer, nv_auth->buffer, nv_auth->size);
    TSS2L_SYS_AUTH_RESPONSE rsp_auths = {0};

    rc = Tss2_Sys_NV_Write(sys_context, auth_handle, nv_index, &cmd_auths,
                           nv_data_write, 0, &rsp_auths);
    return rc;
}

TSS2_RC nv_read_public(TSS2_SYS_CONTEXT *sys_context, TPMI_RH_NV_INDEX nv_index,
                       TPM2B_NV_PUBLIC *nv_public, TPM2B_NAME *nv_name) {
    TSS2_RC rc;

    rc = Tss2_Sys_NV_ReadPublic(sys_context, nv_index, NULL, nv_public, nv_name,
                                NULL);

    return rc;
}

TSS2_RC nv_write_lock(TSS2_SYS_CONTEXT *sys_context, TPMI_RH_NV_INDEX nv_index,
                            TPMI_RH_NV_AUTH auth_handle, TPM2B_AUTH *nv_auth) {
    TSS2_RC rc;
    TSS2L_SYS_AUTH_COMMAND cmd_auths = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = nv_auth->size},
        }
    };
    memcpy(cmd_auths.auths[0].hmac.buffer, nv_auth->buffer, nv_auth->size);
    TSS2L_SYS_AUTH_RESPONSE rsp_auths = {0};

    rc = Tss2_Sys_NV_WriteLock(sys_context, auth_handle, nv_index, &cmd_auths,
                               &rsp_auths);
    return rc;
}

TSS2_RC nv_remove(TSS2_SYS_CONTEXT *sys_context, TPM2B_AUTH *auth,
                  TPMI_RH_NV_INDEX nv_index, TPMI_RH_NV_AUTH auth_handle) {
    TSS2_RC rc;
    TSS2L_SYS_AUTH_COMMAND cmd_auths = {.count = 1,
                                        .auths[0] = {
                                            .sessionHandle = TPM2_RH_PW,
                                            .hmac = {.size = auth->size},
                                        }};
    memcpy(cmd_auths.auths[0].hmac.buffer, auth->buffer, auth->size);
    TSS2L_SYS_AUTH_RESPONSE rsp_auths = {0};

    rc = Tss2_Sys_NV_UndefineSpace(sys_context, auth_handle, nv_index,
                                   &cmd_auths, &rsp_auths);
    return rc;
}

void test_nv_owner(TSS2_SYS_CONTEXT *sys_context) {
    TSS2_RC rc;

    // Parameters of NV Index
    size_t data_size = 32;
    TPM2B_MAX_NV_BUFFER nv_data_read;
    TPM2B_MAX_NV_BUFFER nv_data_write;
    nv_data_write.size = data_size;
    for (size_t i = 0; i < nv_data_write.size; i++)
        nv_data_write.buffer[i] = i + 1;

    TPM2B_AUTH hi_auth = {.size = 0};

    TPM2B_AUTH nv_auth = {.size = 20};
    for (size_t i = 0; i < nv_auth.size; i++)
        nv_auth.buffer[i] = i;

    // Define NV Index
    TPMA_NV attributes = attributes = TPMA_NV_AUTHREAD | TPMA_NV_AUTHWRITE |
                                      TPMA_NV_WRITE_STCLEAR | TPMA_NV_ORDERLY;
    rc = nv_define(sys_context, attributes, TPM2_RH_OWNER, &hi_auth,
                   NV_INDEX_TEST, &nv_auth, data_size);
    CHECK_PASSED(rc);

    // Read NV Index Public
    TPM2B_NAME nv_name;
    TPM2B_NV_PUBLIC nv_public = {0};
    INIT_SIMPLE_TPM2B_SIZE(nv_name);
    rc = nv_read_public(sys_context, NV_INDEX_TEST, &nv_public, &nv_name);
    CHECK_PASSED(rc);

    // Read NV Index
    INIT_SIMPLE_TPM2B_SIZE(nv_data_read);
    rc = nv_read(sys_context, NV_INDEX_TEST, NV_INDEX_TEST, &nv_auth,
                 &nv_data_read, data_size);
    CHECK_FAILED(rc, TPM2_RC_NV_UNINITIALIZED);

    // Write NV Index
    rc = nv_write(sys_context, NV_INDEX_TEST, NV_INDEX_TEST, &nv_auth,
                  &nv_data_write);
    CHECK_PASSED(rc);

    // Read NV Index
    INIT_SIMPLE_TPM2B_SIZE(nv_data_read);
    rc = nv_read(sys_context, NV_INDEX_TEST, NV_INDEX_TEST, &nv_auth,
                 &nv_data_read, data_size);
    CHECK_PASSED(rc);

    // Remove NV Index
    rc = nv_remove(sys_context, &hi_auth, NV_INDEX_TEST, TPM2_RH_OWNER);
    CHECK_PASSED(rc);
}

void test_nv_platform(TSS2_SYS_CONTEXT *sys_context) {
    TSS2_RC rc;

    // Parameters of NV Index
    size_t data_size = 32;
    TPM2B_MAX_NV_BUFFER nv_data_read;
    TPM2B_MAX_NV_BUFFER nv_data_write;
    nv_data_write.size = data_size;
    for (size_t i = 0; i < nv_data_write.size; i++)
        nv_data_write.buffer[i] = i + 1;

    TPM2B_AUTH hi_auth = {.size = 0};
    TPM2B_AUTH nv_auth = {.size = 20};
    for (size_t i = 0; i < nv_auth.size; i++)
        nv_auth.buffer[i] = i;

    TPMA_NV attributes = TPMA_NV_PPREAD | TPMA_NV_PPWRITE |
                         TPMA_NV_WRITE_STCLEAR | TPMA_NV_PLATFORMCREATE |
                         TPMA_NV_ORDERLY;

    // Read NV Index, Failed
    INIT_SIMPLE_TPM2B_SIZE(nv_data_read);
    rc = nv_read(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM, &nv_auth,
                 &nv_data_read, data_size);
    CHECK_FAILED(rc, TPM2_RC_2 + TPM2_RC_HANDLE);

    // Define NV Index
    rc = nv_define(sys_context, attributes, TPM2_RH_PLATFORM, &hi_auth,
                   NV_INDEX_TEST, &nv_auth, data_size);
    CHECK_PASSED(rc);

    // Read NV Index Public
    TPM2B_NAME nv_name;
    INIT_SIMPLE_TPM2B_SIZE(nv_name);
    TPM2B_NV_PUBLIC nv_public = {0};
    rc = nv_read_public(sys_context, NV_INDEX_TEST, &nv_public, &nv_name);
    CHECK_PASSED(rc);

    // Read NV Index, Failed
    INIT_SIMPLE_TPM2B_SIZE(nv_data_read);
    rc = nv_read(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM, &hi_auth,
                 &nv_data_read, data_size);
    CHECK_FAILED(rc, TPM2_RC_NV_UNINITIALIZED);

    // Write NV Index
    rc = nv_write(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM, &hi_auth,
                  &nv_data_write);
    CHECK_PASSED(rc);

    // Read NV Index
    INIT_SIMPLE_TPM2B_SIZE(nv_data_read);
    rc = nv_read(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM, &hi_auth,
                 &nv_data_read, data_size);
    CHECK_PASSED(rc);

    // Re-Define NV Index, Failed
    rc = nv_define(sys_context, attributes, TPM2_RH_PLATFORM, &hi_auth,
                   NV_INDEX_TEST, &nv_auth, data_size);
    CHECK_FAILED(rc, TPM2_RC_NV_DEFINED);

    // Write NV Index
    rc = nv_write(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM, &hi_auth,
                  &nv_data_write);
    CHECK_PASSED(rc);

    // Read NV Index
    INIT_SIMPLE_TPM2B_SIZE(nv_data_read);
    rc = nv_read(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM, &hi_auth,
                 &nv_data_read, data_size);
    CHECK_PASSED(rc);

    // Write and Lock NV Index
    rc = nv_write_lock(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM,
                             &hi_auth);
    CHECK_PASSED(rc);

    // Write NV Index, Failed
    rc = nv_write(sys_context, NV_INDEX_TEST, TPM2_RH_PLATFORM, &hi_auth,
                  &nv_data_write);
    CHECK_FAILED(rc, TPM2_RC_NV_LOCKED);

    // Remove NV Index
    rc = nv_remove(sys_context, &hi_auth, NV_INDEX_TEST, TPM2_RH_PLATFORM);
    CHECK_PASSED(rc);

    // Define NV Index again
    rc = nv_define(sys_context, attributes, TPM2_RH_PLATFORM, &hi_auth,
                   NV_INDEX_TEST, &nv_auth, data_size);
    CHECK_PASSED(rc);

    // Remove NV Index again
    rc = nv_remove(sys_context, &hi_auth, NV_INDEX_TEST, TPM2_RH_PLATFORM);
    CHECK_PASSED(rc);
}

int main(int argc, char *argv[]) {

    TSS2_RC rc;

    // Initialize TCTI contexts
    rc = init_tcti_context(&tcti_context);
    CHECK_RC(rc);

    // Initialize SAPI context
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_RC(rc);

    // Startup TPM
    rc = startup_tpm(sys_context);
    CHECK_RC(rc);

    // Test NV Index [Owner Hierarchy]
    test_nv_owner(sys_context);

    // Test NV Index [Platform Hierarchy]
    test_nv_platform(sys_context);

    // Cleanup resources
    cleanup(rc);

    return 0;
}
